#include <buola/cv/registration/cposegraph.h>

namespace buola { namespace cvi {

struct CNodeInfo
{
    CNodeInfo() {}
    
    CNodeInfo(CPoseGraph::SVertex *pNode,double pTW,double pRW,int pDir,const mat::C3DRotTrans_d &pTransform)
        :   mNode(pNode)
        ,   mTranslationalWeight(pTW)
        ,   mRotationalWeight(pRW)
        ,   mDirection(pDir)
        ,   mTransform(pTransform)
    {}
    
    CPoseGraph::SVertex *mNode;
    
    double mTranslationalWeight;
    double mRotationalWeight;
    int mDirection;
    
    mat::C3DRotTrans_d mTransform;
};

void compute_preconditioner(const CPoseGraph &pGraph,mat::CMat_d &M,mat::CVec2d &pGamma)
{
    M=mat::zeros(pGraph.VertexCount()+1,2);
    
    pGamma[0]=pGamma[1]=std::numeric_limits<double>::max();

    int lEdgeCount=0;
    for(auto e : pGraph.SortedEdges())
    {
        lEdgeCount++;

        mat::CMat<double,6,6> W=e->mInfo;

        CPoseGraph::SVertex *lTop=e->mTop;
        
        for(int dir=0;dir<2;dir++)
        {
            CPoseGraph::SVertex *n=(dir==0)?e->m1:e->m2;
            while(n!=lTop)
            {
                uint i=n->mID;
                double lRW=min_n(W(0,0),W(1,1),W(2,2));
                double lTW=min_n(W(3,3),W(4,4),W(5,5));
                
                M(i,0)+=lRW;
                M(i,1)+=lTW;
                pGamma[0]=min_n(lRW,pGamma[0]);
                pGamma[1]=min_n(lTW,pGamma[1]);
                n=n->mParent;
            }
        }
    }
}


void recompute_transformations(CPoseGraph::SVertex *pV,CPoseGraph::SVertex* pTop)
{
    if(pV==pTop)
        return;
    recompute_transformations(pV->mParent,pTop);
    pV->mTransform=pV->mParent->mTransform*pV->mParameters;
}

void recompute_parameters(CPoseGraph::SVertex *pV,CPoseGraph::SVertex *pTop)
{
    while(pV!=pTop)
    {
        pV->mParameters=inverse(pV->mParent->mTransform)*pV->mTransform;
        pV=pV->mParent;
    }
}

mat::CQuat_d get_rotation(CPoseGraph::SVertex *pV,CPoseGraph::SVertex *pTop)
{
    if (pV==pTop)
        return pV->mTransform.Rotation();
    
    mat::CQuat_d lResult;
    
    while(pV!=pTop)
    {
        lResult=pV->mParameters.Rotation()*lResult;
        pV=pV->mParent;
    }
    return pTop->mTransform.Rotation()*lResult;
}

static const double cTransGain=1.0;
static const double cRotGain=1.0;

void optimize_tree(const CPoseGraph &pGraph)
{
    mat::CMat_d M;
    mat::CVec2d lGamma;

    compute_preconditioner(pGraph,M,lGamma);
    
    std::vector<CNodeInfo> lPath;
    
    for(int lIter=0;lIter<100;lIter++)
    {
        //iterate here
        bool lUsePreconditioner=true;
        int lEdgeCount=0;

        // this is the workspace for computing the paths without 
        // bothering too much the memory allocation
        
        lPath.resize(pGraph.EdgeCount()+1);

        for(CPoseGraph::SEdge *e : pGraph.SortedEdges())
        {
            lEdgeCount++;
            if(!(lEdgeCount%1000))
                msg_info() << "c\n";

            CPoseGraph::SVertex* lTop=e->mTop;
            CPoseGraph::SVertex* l1=e->m1;
            CPoseGraph::SVertex* l2=e->m2;
            int lLength=e->mLength;
    
            recompute_transformations(l1,lTop);
            recompute_transformations(l2,lTop);

            //BEGIN: Path and weight computation 
            int lPC=0;
            CPoseGraph::SVertex* lAux=l1;
            double lTotTW=0,lTotRW=0;
            
            while(lAux!=lTop)
            {
                int lIndex=lAux->mID;
                double lTW=1.0/lLength;
                double lRW=1.0/lLength;

                if (lUsePreconditioner)
                {
                    lTW=1.0/M(lIndex,0); 
                    lRW=1.0/M(lIndex,1);
                }
                lTotTW+=lTW;
                lTotRW+=lRW;
                
                lPath[lPC++]=CNodeInfo(lAux,lTW,lRW,-1,lAux->mTransform);
                lAux=lAux->mParent;
            }
    
            int lTopIndex=lPC;
            lPath[lPC++]=CNodeInfo(lTop,0,0,0,lTop->mTransform);
            
            lPC=lLength;
            lAux=l2;
    
            while(lAux!=lTop)
            {
                int lIndex=lAux->mID;
                double lTW=1.0/lLength;
                double lRW=1.0/lLength;
                if (lUsePreconditioner)
                {
                    lTW=1.0/M(lIndex,0); 
                    lRW=1.0/M(lIndex,1);
                }
                lTotTW+=lTW;
                lTotRW+=lRW;

                lPath[lPC--]=CNodeInfo(lAux,lTW,lRW,-1,lAux->mTransform);
                lAux=lAux->mParent;
            }
            //END: Path and weight computation
    
            //BEGIN: Rotational Error
            mat::CQuat_d lR1=get_rotation(l1,lTop);
            mat::CQuat_d lR2=get_rotation(l2,lTop);
            mat::CQuat_d lRE=e->mTransform.Rotation();
            mat::CQuat_d lRR=conj(lR2)*lR1*lRE;

            double lRotationFactor=lUsePreconditioner?sqrt(double(lLength))*min_n(e->mInfo(0,0),e->mInfo(1,1),e->mInfo(2,2))/
                                                                                                (lGamma[0]*(double)lIter)
                                                     :sqrt(double(lLength))*cRotGain/(double)lIter;

            if(lRotationFactor>1) lRotationFactor=1;

            mat::CQuat_d lTotalRotation=lPath[lLength].mTransform.Rotation()*lRR*conj(lPath[lLength].mTransform.Rotation());
            mat::CVec3d lAxis=lTotalRotation.GetAxis();
            double lAngle=arg(lTotalRotation);

            double lCW=0;
            for(int i=1;i<=lTopIndex;i++)
            {
                lCW+=lPath[i-1].mRotationalWeight/lTotRW;
                mat::CQuat_d lR=lPath[i].mTransform.Rotation();
                mat::CQuat_d lB(lAngle*lCW*lRotationFactor,lAxis);
                lPath[i].mTransform.Rotation()=lB*lR;
            }
            for(int i=lTopIndex+1;i<=lLength;i++)
            {
                lCW+=lPath[i].mRotationalWeight/lTotRW;
                mat::CQuat_d lR=lPath[i].mTransform.Rotation();
                mat::CQuat_d lB(lAngle*lCW*lRotationFactor,lAxis);
                lPath[i].mTransform.Rotation()=lB*lR;
            }

            //recompute the parameters based on the transformation
            for(int i=0;i<lTopIndex;i++)
            {
                CPoseGraph::SVertex* lNode=lPath[i].mNode;
                lNode->mParameters.Rotation()=conj(lPath[i+1].mTransform.Rotation())*lPath[i].mTransform.Rotation();
            }
            for (int i=lTopIndex+1;i<=lLength;i++)
            {
                CPoseGraph::SVertex* lNode=lPath[i].mNode;
                lNode->mParameters.Rotation()=conj(lPath[i-1].mTransform.Rotation())*lPath[i].mTransform.Rotation();
            }
            //END: Rotational Error

            //now spread the parameters
            recompute_transformations(l1,lTop);
            recompute_transformations(l2,lTop);

            //BEGIN: Translational Error
            //CHECK THIS!!!!!!!!!!!!!!!!!
            mat::CVec3d lTR=l1->mTransform.Rotation()*e->mTransform.Translation()+
                            l1->mTransform.Translation()-l2->mTransform.Translation();

            double lTranslationFactor=lUsePreconditioner?cTransGain*lLength*min_n(e->mInfo(3,3),e->mInfo(4,4),e->mInfo(5,5))/
                                                                        (lGamma[1]*(double)lIter)
                                                        :cTransGain*lLength/(double)lIter;

            if(lTranslationFactor>1) lTranslationFactor=1;
            
            mat::CVec3d lDT=lTR*lTranslationFactor;

            //left wing
            double lLCum=0;
            for (int i=lTopIndex-1;i>=0;i--)
            {
                CPoseGraph::SVertex* n=lPath[i].mNode;
                lLCum-=(lUsePreconditioner)?lPath[i].mTranslationalWeight/lTotTW:1.0/lLength;
                n->mTransform.Translation()+=lLCum*lDT;
            }
            //right wing
            double lRCum=0;
            for (int i=lTopIndex+1; i<=lLength;i++)
            {
                CPoseGraph::SVertex* n=lPath[i].mNode;
                lRCum+=(lUsePreconditioner)?lPath[i].mTranslationalWeight/lTotTW:1.0/lLength;
                n->mTransform.Translation()+=lRCum*lDT;
            }
            
            assert(abs(lLCum+lRCum)-1<1e-6);

            recompute_parameters(l1,lTop);
            recompute_parameters(l2,lTop);
            //END: Translational Error
        }
        
//         if(pAdaptiveRestart)
//         {
//             double mte, ate;
//             double mre, are;
//             error(&mre, &mte, &are, &ate);
//             maxTranslationalErrors.push_back(mte);
//             maxRotationalErrors.push_back(mre);
//             int interval=3;
//             if ((int)maxRotationalErrors.size()>=interval)
//             {
//                 uint s=maxRotationalErrors.size();
//                 double re0 = maxRotationalErrors[s-interval];
//                 double re1 = maxRotationalErrors[s-1];
// 
//                 if ((re1-re0)>are || sqrt(re1)>0.99*M_PI)
//                 {
//                     double rg=rotGain;
//                     if (sqrt(re1)>M_PI/4)
//                     {
//                         cerr << "RESTART!!!!! : Angular wraparound may be occourring" << endl;
//                         cerr << " err=" << re0 << " -> " << re1 << endl; 
//                         cerr << "Restarting optimization and reducing the rotation factor" << endl;
//                         cerr << rg << " -> ";
//                         initializeOnTree();
//                         initializeTreeParameters();
//                         initializeOptimization();
//                         error(&mre, &mte);
//                         maxTranslationalErrors.push_back(mte);
//                         maxRotationalErrors.push_back(mre);
//                         rg*=0.1;
//                         rotGain=rg;
//                         cerr << rotGain << endl;
//                     }
//                     else 
//                     {
//                         cerr << "decreasing angular gain" << rotGain*0.1 << endl;
//                         rotGain*=0.1;
//                     }
//                 }
//             }
//         }
    }
}
    
/*namespace cvi*/ } /*namespace buola*/ }

#if 0

/**********************************************************************
 *
 * This source code is part of the Tree-based Network Optimizer (TORO)
 *
 * TORO Copyright (c) 2007 Giorgio Grisetti, Cyrill Stachniss, 
 *                         Slawomir Grzonka, and Wolfram Burgard
 *
 * TORO is licences under the Common Creative License,
 * Attribution-NonCommercial-ShareAlike 3.0
 *
 * You are free:
 *   - to Share - to copy, distribute and transmit the work
 *   - to Remix - to adapt the work
 *
 * Under the following conditions:
 *
 *   - Attribution. You must attribute the work in the manner specified
 *     by the author or licensor (but not in any way that suggests that
 *     they endorse you or your use of the work).
 *  
 *   - Noncommercial. You may not use this work for commercial purposes.
 *  
 *   - Share Alike. If you alter, transform, or build upon this work,
 *     you may distribute the resulting work only under the same or
 *     similar license to this one.
 *
 * Any of the above conditions can be waived if you get permission
 * from the copyright holder.  Nothing in this license impairs or
 * restricts the author's moral rights.
 *
 * TORO is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied 
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
 * PURPOSE.  
 **********************************************************************/

/** \file treeoptimizer3.cpp
 *
 * \brief Defines the core optimizer class for 3D graphs which is a
 * subclass of TreePoseGraph3
 *
 **/

#include "treeoptimizer3.hh"
#include <string>

using namespace std;

namespace AISNavigation {

#define DEBUG(i) \
        if (verboseLevel>i) cerr


TreeOptimizer3::TreeOptimizer3(){
  restartOnDivergence=false;
  sortedEdges=0;
  mpl=-1;
  edgeCompareMode=EVComparator<Edge*>::CompareLevel;
}


TreeOptimizer3::~TreeOptimizer3(){
}

void TreeOptimizer3::initializeTreeParameters(){
  ParameterPropagator pp;
  treeDepthVisit(pp,root);
}



void TreeOptimizer3::recomputeTransformations(Vertex*v, Vertex* top){
  if (v==top)
    return;
  recomputeTransformations(v->parent, top);
  v->transformation=v->parent->transformation*v->parameters;
}


void TreeOptimizer3::recomputeParameters(Vertex*v, Vertex* top){
  while (v!=top){
    v->parameters=v->parent->transformation.inv()*v->transformation;
    v=v->parent;
  }
}


TreeOptimizer3::Transformation TreeOptimizer3::getPose(Vertex*v, Vertex* top){
  Transformation t(0.,0.,0.,0.,0.,0.);
  if (v==top)
    return v->transformation;
  while (v!=top){
    t=v->parameters*t;
    v=v->parent;
  }
  return top->transformation*t;
}

TreeOptimizer3::Rotation TreeOptimizer3::getRotation(Vertex*v, Vertex* top){
  Rotation r(0.,0.,0.);
  if (v==top)
    return v->transformation.rotation();
  while (v!=top){
    r=v->parameters.rotation()*r;
    v=v->parent;
  }
  return top->transformation.rotation()*r;
}



double TreeOptimizer3::error(const Edge* e) const{
  const Vertex* v1=e->v1;
  const Vertex* v2=e->v2;
  
  Transformation et=e->transformation;
  Transformation t1=v1->transformation;
  Transformation t2=v2->transformation;
  
  Transformation t12=(t1*et)*t2.inv();
  Pose p12=t12.toPoseType();

  Pose ps=e->informationMatrix*p12;
  double err=p12*ps;
  DEBUG(100) << "e(" << v1->id << "," << v2->id << ")" << err << endl;
  return err;
 
}

double TreeOptimizer3::traslationalError(const Edge* e) const{
  const Vertex* v1=e->v1;
  const Vertex* v2=e->v2;
  
  Transformation et=e->transformation;
  Transformation t1=v1->transformation;
  Transformation t2=v2->transformation;
  
  Translation t12=(t2.inv()*(t1*et)).translation();
  return t12*t12;;
 
}

double TreeOptimizer3::rotationalError(const Edge* e) const{
  const Vertex* v1=e->v1;
  const Vertex* v2=e->v2;
  
  Rotation er=e->transformation.rotation();
  Rotation r1=v1->transformation.rotation();
  Rotation r2=v2->transformation.rotation();
  
  Rotation r12=r2.inverse()*(r1*er);
  double r=r12.angle();
  return r*r;
}


double TreeOptimizer3::loopError(const Edge* e) const{
  double err=0;
  const Vertex* v=e->v1;
  while (v!=e->top){
    err+=error(v->parentEdge);
    v=v->parent;
  }
  v=e->v2;
  while (v==e->top){
    err+=error(v->parentEdge);
    v=v->parent;
  }
  if (e->v2->parentEdge!=e && e->v1->parentEdge!=e)
    err+=error(e);
  return err;
}

double TreeOptimizer3::loopRotationalError(const Edge* e) const{
  double err=0;
  const Vertex* v=e->v1;
  while (v!=e->top){
    err+=rotationalError(v->parentEdge);
    v=v->parent;
  }
  v=e->v2;
  while (v!=e->top){
    err+=rotationalError(v->parentEdge);
    v=v->parent;
  }
  if (e->v2->parentEdge!=e && e->v1->parentEdge!=e)
    err+=rotationalError(e);
  return err;
}


double TreeOptimizer3::error(double* mre, double* mte, double* are, double* ate, TreePoseGraph3::EdgeSet* eset) const{
  double globalRotError=0.;
  double maxRotError=0;
  double globalTrasError=0.;
  double maxTrasError=0;

  int c=0;
  if (! eset){
    for (TreePoseGraph3::EdgeMap::const_iterator it=edges.begin(); it!=edges.end(); it++){
      double re=rotationalError(it->second);
      globalRotError+=re;
      maxRotError=maxRotError>re?maxRotError:re;
      double te=traslationalError(it->second);
      globalTrasError+=te;
      maxTrasError=maxTrasError>te?maxTrasError:te;
      c++;
    }
  } else {
    for (TreePoseGraph3::EdgeSet::const_iterator it=eset->begin(); it!=eset->end(); it++){
      const TreePoseGraph3::Edge* edge=*it;
      double re=rotationalError(edge);
      globalRotError+=re;
      maxRotError=maxRotError>re?maxRotError:re;
      double te=traslationalError(edge);
      globalTrasError+=te;
      maxTrasError=maxTrasError>te?maxTrasError:te;
      c++;
    }
  }

  if (mte)
    *mte=maxTrasError;
  if (mre)
    *mre=maxRotError;
  if (ate)
    *ate=globalTrasError/c;
  if (are)
    *are=globalRotError/c;
  return globalRotError+globalTrasError;
}



void TreeOptimizer3::initializeOptimization(EdgeCompareMode mode){
  edgeCompareMode=mode;
  // compute the size of the preconditioning matrix
  int sz=maxIndex()+1;
  DEBUG(1) << "Size= " << sz << endl;
  M.resize(sz);

  DEBUG(1) << "allocating M(" << sz << ")" << endl;
  iteration=1;

  // sorting edges
  if (sortedEdges!=0){
    delete sortedEdges;
    sortedEdges=0;
  }
  sortedEdges=sortEdges();
  mpl=maxPathLength();
  rotGain=1.;
  trasGain=1.;
}

void TreeOptimizer3::initializeOnlineIterations(){
  int sz=maxIndex()+1;
  DEBUG(1) << "Size= " << sz << endl;
  M.resize(sz);
  DEBUG(1) << "allocating M(" << sz << ")" << endl;
  iteration=1;
  maxRotationalErrors.clear();
  maxTranslationalErrors.clear();
  rotGain=1.;
  trasGain=1.;
}

void TreeOptimizer3::initializeOnlineOptimization(EdgeCompareMode mode){
  edgeCompareMode=mode;
  // compute the size of the preconditioning matrix
  clear();
  Vertex* v0=addVertex(0,Pose(0,0,0,0,0,0));
  root=v0;
  v0->parameters=Transformation(v0->pose);
  v0->parentEdge=0;
  v0->parent=0;
  v0->level=0;
  v0->transformation=Transformation(TreePoseGraph3::Pose(0,0,0,0,0,0));
}

}; //namespace AISNavigation




/**********************************************************************
 *
 * This source code is part of the Tree-based Network Optimizer (TORO)
 *
 * TORO Copyright (c) 2007 Giorgio Grisetti, Cyrill Stachniss, 
 *                         Slawomir Grzonka and Wolfram Burgard
 *
 * TORO is licences under the Common Creative License,
 * Attribution-NonCommercial-ShareAlike 3.0
 *
 * You are free:
 *   - to Share - to copy, distribute and transmit the work
 *   - to Remix - to adapt the work
 *
 * Under the following conditions:
 *
 *   - Attribution. You must attribute the work in the manner specified
 *     by the author or licensor (but not in any way that suggests that
 *     they endorse you or your use of the work).
 *  
 *   - Noncommercial. You may not use this work for commercial purposes.
 *  
 *   - Share Alike. If you alter, transform, or build upon this work,
 *     you may distribute the resulting work only under the same or
 *     similar license to this one.
 *
 * Any of the above conditions can be waived if you get permission
 * from the copyright holder.  Nothing in this license impairs or
 * restricts the author's moral rights.
 *
 * TORO is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied 
 * warranty of MERCHANTABILITY or FITNESS FOR A PARTICULAR
 * PURPOSE.  
 **********************************************************************/

#include "treeoptimizer3.hh"
#include <string>

using namespace std;

namespace AISNavigation {

#define DEBUG(i) \
        if (verboseLevel>i) cerr



//helper functions. Should I explain :-)?
inline double max3( const double& a, const double& b, const double& c){
  double m=a>b?a:b;
  return m>c?m:c;
}
inline double min3( const double& a, const double& b, const double& c){
  double m=a<b?a:b;
  return m<c?m:c;
}


struct NodeInfo{
  TreeOptimizer3::Vertex* n;
  double translationalWeight;
  double rotationalWeight;
  int direction;
  TreeOptimizer3::Transformation transformation;
  TreeOptimizer3::Transformation parameters;
  NodeInfo(TreeOptimizer3::Vertex* v=0, double tw=0, double rw=0, int dir=0,
       TreeOptimizer3::Transformation t=TreeOptimizer3::Transformation(0,0,0,0,0,0), 
       TreeOptimizer3::Parameters p=TreeOptimizer3::Transformation(0,0,0,0,0,0)){
    n=v;
    translationalWeight=tw;
    rotationalWeight=rw;
    direction=dir;
    transformation=t;
    parameters=p;
  }
};

typedef std::vector<NodeInfo> NodeInfoVector;


/********************************** Preconditioned and unpreconditioned error distribution ************************************/





void TreeOptimizer3::computePreconditioner(){
  for (uint i=0; i<M.size(); i++){
    M[i][0]=0;
    M[i][1]=0;
  }
  gamma[0] = gamma[1] = numeric_limits<double>::max();

  int edgeCount=0;
  for (EdgeSet::iterator it=sortedEdges->begin(); it!=sortedEdges->end(); it++){
    edgeCount++;
    if (! (edgeCount%1000))
      DEBUG(1) << "m";

    Edge* e=*it;
    Transformation t=e->transformation;
    InformationMatrix W=e->informationMatrix;
    
    Vertex* top=e->top;
    for (int dir=0; dir<2; dir++){
      Vertex* n = (dir==0)? e->v1 : e->v2;
      while (n!=top){
    uint i=n->id;
    double rW=min3(W[0][0], W[1][1], W[2][2]);
    double tW=min3(W[3][3], W[4][4], W[5][5]);
    M[i][0]+=rW;
    M[i][1]+=tW;
    gamma[0]=gamma[0]<rW?gamma[0]:rW;
    gamma[1]=gamma[1]<tW?gamma[1]:tW;
    n=n->parent;
      }
    }

  }
  
  if (verboseLevel>1){
    for (uint i=0; i<M.size(); i++){
      cerr << "M[" << i << "]=" << M[i][0] << " " << M[i][1] << endl;
    }
  }
}



};//namespace AISNavigation



#endif